import os
import get_lines as gl
from multiprocessing import Pool


def get_stocks(year, stocks):
    gl.login()
    if not os.path.exists('stocks/%d' % year):
        os.mkdir('stocks/%d' % year)
        print('文件夹创建成功: ', year)
    try:
        for index, next_data in stocks.iterrows():
            code = next_data['code']
            save_path = 'stocks/%d/%s.csv' % (year, code)
            if os.path.exists(save_path):
                continue
            print('开始 %s, %d' % (code, year))
            stock = gl.query_history_k_data_plus(code, '%d-01-01' % year, '%d-01-01' % (year + 1))
            print('结束 %s, %d' % (code, year))
            if stock.empty:
                continue
            stock['name'] = next_data['code_name']
            stock.to_csv(save_path)
    except Exception as e:
        print(e)
    gl.logout()


if __name__ == '__main__':
    pool = Pool(8)
    gl.login()
    sto = gl.get_stocks().get_data()
    for y in range(2013, 2021):
        pool.apply_async(get_stocks, (y, sto))
    pool.close()
    pool.join()
    gl.logout()
